Emit affine maps from AutoTP layers - #8519

Step 4 of the staging plan in #8252: AutoTP layers now emit an affine map, and the converter reads it on a path a real job takes rather than one a test supplies.

Follows #8385 (the IR, the lowering constructors, the converter branch) and #8477 (the scale fix). Scope here is contiguous splits and replicated parameters; the fused QKV and Yuan shared-QK layouts are next and still carry unsupported_reason, so they emit no map and are unaffected.

Where a map is built, and why it is not where you would expect

The producer cannot build a map from what collect_autotp_universal_checkpoint_info sees. The conversion metadata carries eight fields and per-rank partition sizes are not among them — those were deliberately restore-only. But the layer does have them: _freeze_partition_sizes resolves them via get_shard_size_list while the layer is built, and they are not recoverable later from a shape alone.

So the map is derived at mark time, in _set_param_uc_meta, which already receives partition_sizes, sub_param_shard_widths, logical_shape and replicated, and has tp_world_size in scope. That puts the derivation in one place rather than in each of the seven _mark_uc_metadata implementations, and collect_... then gathers what the layers produced.

A layout that is not describable yet returns None, so conversion falls back to the categories.

The map and the categories have to coexist

A checkpoint carrying an affine map also carries the existing pattern lists, so a converter predating the map still reads it — that is the additivity §6.3 promises.

The consequence is that the converter which prefers the map never consults those branches, which leaves their patterns looking unused and fails a strict conversion. They are superseded, not unused, so the affine branch now marks them consumed. This is the compatibility question raised as "one boundary for review" in #8385; it turns out both have to be present and the reader has to account for the other.

This only surfaced end to end. The producer's output was correct in isolation, and the emitted maps matched a layout already verified by a passing resume — the failure was a strict-mode assert on rank 0 during conversion, which presented as rank 1 blocking on the following barrier.

Tests

TestAffineMapProducer builds a real AutoTP engine at TP2 and requires the producer to emit exactly the four maps the resume fixture previously supplied by hand. Those maps are not a guess: a full train → save → convert → resume cycle reproduces uninterrupted training through them, so matching them is evidence rather than self-consistency.

affine_resume_checkpoint now takes its layout from the producer instead of injecting one, so the four save-convert-resume cases exercise the metadata a real job writes. All four pass — TP2 → TP1 and TP2 → TP2, through both the legacy and affine paths, compared against uninterrupted training on logits, losses, gradients, FP32 weights, both Adam moments and step counters.

Validated on CPU/gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4): 109 passed across the resume matrix, the producer and coverage tests, the vocab cross-TP test, the affine unit suite and tests/unit/runtime/tensor_parallel/. The five remaining failures in that file are FusedAdam JIT-compile errors on this machine and are identical on upstream master.

Review findings

Three things came out of review, all fixed here:
- affine_map was added to every conversion metadata dict, changing the schema even for layouts with no map. Now emitted only where there is one, so an existing layer publishes exactly what it did before.
- The producer test checked that the expected maps were present but not that nothing extra was. It now asserts the pattern set exactly.
- Parameters the TP machinery leaves untouched — embeddings, norm weights — were classified TP_REPLICATED with no map, so conversion still fell back to the category. TestAffineMapCoverage now asserts the general invariant: every pattern in any category must carry a map, with only AutoTP-unsupported layouts exempt. It fails without the fix.
- TestUnevenVocabCrossTpResume covers a 101-row vocabulary head saved at TP2 and restored at TP1 and TP2. The 51/50 split is the point — restoring at TP1 merges two unequal shards, where an even-split assumption would surface. The fixture asserts the emitted map records that split before saving. Adapted from @jinyouzhi's cross-tp test in #8309, using the column partition already in tree so it carries no dependency on that PR.

affine_ir_spec.md gains §6.4 on where a map is built, and a §6.3 rule on the map and categories coexisting.

Related: #8252, #8230.

cc @delock